#!/usr/bin/env python3
'''
未验证
'''
import gymnasium as gym
import argparse

from PIL import Image
from lib import common

import numpy as np
import torch
import torch.nn as nn
import ptan

from typing import Any
import ale_py

gym.register_envs(ale_py)

class AtariA2C(nn.Module):
    def __init__(self, obs_size, n_actions):
        super(AtariA2C, self).__init__()

        self.encoder = nn.Sequential(
            nn.Linear(obs_size, 1024),
            nn.ReLU(),
            nn.Linear(1024, 1024),
            nn.ReLU(),
        )

        self.policy = nn.Sequential(
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Linear(512, n_actions)
        )

        self.value = nn.Sequential(
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Linear(512, 1)
        )

    def forward(self, x):
        linear_out = self.encoder(x.float() / 255.0)
        return self.policy(linear_out), self.value(linear_out)


class RamFrameStack(gym.Wrapper):
    def __init__(self, env, k):
        gym.Wrapper.__init__(self, env)
        self.k = k
        # 创建一个最大只有K帧的队列，一旦长度超过K，那么自动将队头的帧去除
        self.frames = deque([], maxlen=k)
        shp = env.observation_space.shape
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(shp[0]*k,), dtype=np.float32)

    def reset(self, seed: int | None = None, options: dict[str, Any] | None = None):
        ob, info = self.env.reset(seed=seed, options=options)
        for _ in range(self.k):
            self.frames.append(ob)
        return self._get_ob(), info

    def step(self, action):
        ob, reward, done, truncated, info = self.env.step(action)
        self.frames.append(ob)
        return self._get_ob(), reward, done, truncated, info

    def _get_ob(self):
        assert len(self.frames) == self.k
        return np.concatenate(self.frames)

    

def wrap_dqn(env, stack_frames=4, episodic_life=True, reward_clipping=True):
    if episodic_life:
        # 将多条生命的游戏模拟成单条生命ActorCriticAgent
        env = ptan.common.wrappers.EpisodicLifeEnv(env)
    # 增强初始化
    env = ptan.common.wrappers.NoopResetEnv(env, noop_max=30)

    if 'FIRE' in env.unwrapped.get_action_meanings():
        env = ptan.common.wrappers.FireResetEnv(env)
    env = RamFrameStack(env, stack_frames)
    return env



if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-m", "--model", default=r"M:\Projects\python\my_-nqd\learning\CarRacing\saves\a2c-ppo_carracing_linear\a2c-best-935.19999999999.pth", help="Model file to load")
    parser.add_argument("-r", "--record", help="If specified, sets the recording dir, default=Disabled")
    parser.add_argument("-s", "--save", type=int, help="If specified, save every N-th step as an image")
    parser.add_argument("--cuda", default=True, action="store_true", help="Enable cuda")
    args = parser.parse_args()

    env = wrap_dqn(gym.make("ALE/BankHeist-ram-v5", frameskip=4, repeat_action_probability=0.0, obs_type="ram", render_mode="human"))
    if args.record:
        env = gym.wrappers.Monitor(env, args.record)

    net = AtariA2C(env.observation_space.shape[0], env.action_space.n)
    net.load_state_dict(torch.load(args.model))
    net = net.to(common.select_device(args))

    obs, _ = env.reset()
    total_reward = 0.0
    total_steps = 0
    while True:
        obs_v = torch.FloatTensor([obs]).to("cuda")
        mu_v = net(obs_v)[0].to("cpu")
        action = mu_v.squeeze(dim=0).data.numpy()
        obs, reward, done, trunc, _ = env.step(np.argmax(action))
        done = done or trunc
        total_reward += reward
        total_steps += 1
        if done:
            break
    print("In %d steps we got %.3f reward" % (total_steps, total_reward))
